#!/usr/bin/env python

import os
import sys
from timeit import default_timer as timer
import numpy as np
import onnxruntime as onnxrt
import pdb

from py_lib.logger_utils import logger_init
from py_lib.ort_test.arg_parser import ort_test_args

# onnxrt.set_default_logger_severity(0)

float_dict = {
    'tensor(float16)': 'float16',
    'tensor(float)': 'float32',
    'tensor(double)': 'float64',
}

integer_dict = {
    'tensor(int32)': 'int32',
    'tensor(int8)': 'int8',
    'tensor(uint8)': 'uint8',
    'tensor(int16)': 'int16',
    'tensor(uint16)': 'uint16',
    'tensor(int64)': 'int64',
    'tensor(uint64)': 'uint64',
}


def generate_feeds(sess, symbolic_dims={}):  # noqa: B006
    feeds = {}
    for input_meta in sess.get_inputs():
        # replace any symbolic dimensions
        shape = []
        for dim in input_meta.shape:
            if not dim:
                # unknown dim
                shape.append(1)
            elif type(dim) == str:
                # symbolic dim. see if we have a value otherwise use 1
                if dim in symbolic_dims:
                    shape.append(int(symbolic_dims[dim]))
                else:
                    shape.append(1)
            else:
                shape.append(dim)

        if input_meta.type in float_dict:
            feeds[input_meta.name] = np.random.rand(*shape).astype(float_dict[input_meta.type])
        elif input_meta.type in integer_dict:
            feeds[input_meta.name] = np.random.uniform(high=1000, size=tuple(shape)).astype(
                integer_dict[input_meta.type]
            )
        elif input_meta.type == 'tensor(bool)':
            feeds[input_meta.name] = np.random.randint(2, size=tuple(shape)).astype('bool')
        else:
            print(f'unsupported input type {input_meta.type} for input {input_meta.name}')
            sys.exit(-1)
    return feeds


# simple test program for loading onnx model, feeding all inputs and running the model num_iters times.
def run_model(args):
    model_path = args.model_path
    num_iters = args.num_iters
    profile = args.profile
    symbolic_dims = args.symbolic_dims
    feeds = None
    override_initializers = True
    
    sess_options = None
    if profile:
        sess_options = onnxrt.SessionOptions()
        sess_options.enable_profiling = True
        sess_options.profile_file_prefix = os.path.basename(model_path)

    sess = onnxrt.InferenceSession(
        model_path,
        sess_options=sess_options,
        providers=onnxrt.get_available_providers(),
    )
    meta = sess.get_modelmeta()

    if not feeds:
        feeds = generate_feeds(sess, symbolic_dims)

    if override_initializers:
        # Starting with IR4 some initializers provide default values
        # and can be overridden (available in IR4). For IR < 4 models
        # the list would be empty
        for initializer in sess.get_overridable_initializers():
            shape = [dim if dim else 1 for dim in initializer.shape]
            if initializer.type in float_dict:
                feeds[initializer.name] = np.random.rand(*shape).astype(float_dict[initializer.type])
            elif initializer.type in integer_dict:
                feeds[initializer.name] = np.random.uniform(high=1000, size=tuple(shape)).astype(
                    integer_dict[initializer.type]
                )
            elif initializer.type == 'tensor(bool)':
                feeds[initializer.name] = np.random.randint(2, size=tuple(shape)).astype('bool')
            else:
                print(f'unsupported initializer type {initializer.type} for initializer {initializer.name}')
                sys.exit(-1)

    start = timer()
    for _i in range(num_iters):
        outputs = sess.run([], feeds)  # fetch all outputs
    end = timer()

    print(f'model: {meta.graph_name}')
    print(f'version: {meta.version}')
    print(f'iterations: {num_iters}')
    print(f'avg latency: {((end - start) * 1000) / num_iters} ms')

    if profile:
        trace_file = sess.end_profiling()
        print(f'trace file written to: {trace_file}')

    return 0, feeds, num_iters > 0 and outputs


def main():
    args = ort_test_args()
    exit_code, _, _ = run_model(args)
    sys.exit(exit_code)

if __name__ == '__main__':
    main()
